"""
Author: qinguoming
Date: 2025-02-11 
Version: 1.0 
Description: 
    1.这个程序是一个模板,用于对桁架单元组成的结构的静态分析,输出节点位移,输出单元应力,单元应变,节点支反力等结果.
    2.支持T1D2,T2D2,T3D2三种单元
    3.节点和网格的编号必须从1开始,且连续.
    4.仅支持节点集中载荷加载
脚本执行: 

python D:\FE3d_dev\FE3d\sol\Truss_Static.py -m D:\FE3d_dev\FE3d\test\test_truss_static\2d_truss_mesh.inp -o result.h5 -c D:\FE3d_dev\FE3d\test\test_truss_static\config.json

测试文件:
在qinguoming.xyz

"""
import numpy as np
import sys
import os
# 将FE3d目录加入系统路径,所以其他的import 要用相对导入..,.
p = os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir))
sys.path.append(p)
print(p)
import argparse,json
import h5py
from src.node import Node
from src.Truss import T1D2, T2D2, T3D2
from src.Reader import InpReader
from src.tools import get_ithdof_index,initNumpyPrintOptions
from src.Solver import AddBcsByBigNumber, Solve_Linear_Static,AddBcsByReorderK
# -------------------
initNumpyPrintOptions()  # 初始化numpy打印选项,使得输出更美观
# 创建解析器对象
parser = argparse.ArgumentParser(description='对桁架单元组成的结构进行静态分析')
# 添加参数
parser.add_argument('-m', type=str, required=True, help='网格文件路径')
parser.add_argument('-c', type=str, required=True, help='设置文件路径')
parser.add_argument('-o', type=str, required=False,default='output.h5', help='输出dat文件路径')
truss_type_choices = [T1D2, T2D2, T3D2]
def truss_sections(prop,i):
    """
    根据输入参数A,E,rho,生成对应单元i的材料属性
    
    args:
        prop: dict,包含A,E,rho等单元属性
        i: int,单元编号
    return:
        list,返回对应单元的材料属性
    """
    if isinstance(prop,list):
        return prop[i-1]
    elif isinstance(prop,dict):
        return prop[str(i)]
    elif isinstance(prop,float) or isinstance(prop,int):
        return prop
    return prop

if __name__ == '__main__':
    args = parser.parse_args()
    # 网格数据
    inp_file = args.m
    # 输出文件
    H5_file = args.o
    # 设置文件
    with open(args.c, 'r') as f:
        S=json.load(f)
    E,A,rho= S["MATERIAL"]["E"],S["MATERIAL"]["A"],S["MATERIAL"]["RHO"]
    bc_dict=S["BOUNDARY"]
    load_dict=S["LOAD"]
    bcprocess=S["OPTION"]["ADD_BC"]
    Truss_dim=S["OPTION"]["ELEM_DIMENSIOIN"]  # 桁架单元的维度,1,2,3

    # 读取inp,并转化为0-based的编号规则
    node_data, element_data, nset_data, elset_data ,surfaces = InpReader(inp_file)
    element_data[:,1:]-=1  # 节点和单元编号从0开始
    
    # 初始化节点
    nodes:list[Node]=[]
    for i in range(node_data.shape[0]):
        nodes.append(Node(label=node_data[i,0],
                            x=node_data[i,1],
                            y=node_data[i,2],
                            z=node_data[i,3]))
    # 初始化单元
    elements=[]
    for i in range(element_data.shape[0]):
        elements.append(truss_type_choices[Truss_dim-1](
            label=element_data[i,0],
            node1=nodes[element_data[i,1]],
            node2=nodes[element_data[i,2]],
            E=truss_sections(E,element_data[i,0]),
            A=truss_sections(A,element_data[i,0]),
            rho=truss_sections(rho,element_data[i,0])
            ))
        print(elements[-1])
    
    # define global stiffness matrix and mass matrix
    global_dof_num=node_data.shape[0]*Truss_dim  # 全局自由度数目
    K=np.zeros((global_dof_num,global_dof_num))
    # M=np.zeros((global_dof_num,global_dof_num))
    
    # assembly K,M
    for e in elements:
        K[np.ix_(e.ElemDofIndexs,e.ElemDofIndexs)]+=e.Ke
        # M[np.ix_(e.ElemDofIndexs,e.ElemDofIndexs)]+=e.Me
    print("global Stiffness matrix done!")
    
    # 边界条件
    bc_ind,bc_val=[],[]
    for bc_data in bc_dict.values():
        for nlabel in nset_data[bc_data["NSET"]]:
            bc_ind.append(get_ithdof_index(nlabel,Truss_dim,bc_data["DOF"]))
            bc_val.append(bc_data["VALUE"])
    bc_val=np.array(bc_val)
    bcsInd= np.array(list(set(bc_ind))) # 去重
    bcsVal = np.array([np.sum(bc_val[bc_ind==bcsInd[i]]) for i in range(bcsInd.shape[0])])
    bcs=np.concatenate((bcsInd[:,np.newaxis],bcsVal[:,np.newaxis]),axis=1)  # (n,2),第一列是自由度索引,第二列是边界条件值
    print(f"Boundary conditions: {bcs}")
    # 加载载荷
    load_arr=np.zeros(global_dof_num)
    for load_data in load_dict.values():
        for nindex in nset_data[load_data["NSET"]]:
            dof_ind=get_ithdof_index(nindex,Truss_dim,load_data["DOF"])
            load_arr[dof_ind]+=load_data["VALUE"]
    forces=np.zeros((load_arr.nonzero()[0].shape[0],2))
    forces[:,0]=load_arr.nonzero()[0]  # 获取非零元素的索引 
    forces[:,1]=load_arr[load_arr.nonzero()]  # 获取非零元素的值
    print("Loads done!")
    print(f"forces={forces}")
    
    # introduce the boundary condition & loads to the global stiffness matrix
    if bcprocess==1:
        K_,F_=AddBcsByReorderK(K,bcs,forces)
    elif bcprocess==2:
        K_,F_=AddBcsByBigNumber(K,bcs,forces)
    else:
        raise ValueError("Unknown boundary condition process method!")
    
    # from src.tools import imShowMat
    # imShowMat(K_,"K_")
    # imShowMat(K,"K")
    
    # solve the linear static problem
    U,F=Solve_Linear_Static(K,K_,F_)
    
    # write the results to the output file
    with h5py.File(H5_file,"w") as f:
        # 写入节点数据,(n,3)
        f.create_dataset("Nodes",data=node_data[:,1:]) 
        # 写入单元数据,(n,3)
        f.create_dataset("Cells",data=np.concatenate((np.ones((element_data.shape[0],1),dtype=int)*element_data[:,1:].shape[1],element_data[:,1:]),axis=1),dtype=int)
        # 写入cell type
        f.create_dataset("Cell_types",data=np.ones(element_data.shape[0],dtype=int)*truss_type_choices[Truss_dim-1].VTK_TYPE,dtype=int)
        # write result data
        ResultGroup=f.create_group(name="Result")
        for i in range(Truss_dim):
            ResultGroup.create_dataset(f"U{i+1}",data=U[i::Truss_dim])  # U1,U2,U3
            ResultGroup.create_dataset(f"F{i+1}",data=F[i::Truss_dim])  # F1,F2,F3
        ResultGroup.create_dataset("Umag",data=np.sqrt(sum([U[i::Truss_dim]**2 for i in range(Truss_dim)])))
        ResultGroup.create_dataset("Fmag",data=np.sqrt(sum([F[i::Truss_dim]**2 for i in range(Truss_dim)])))
    
    print(f"Results written to {H5_file} successfully!")
    print("Done!")  
    # 结束脚本